

#include "ju.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <unistd.h> // 包含此行以声明 getpid 函数
#include <sys/uio.h> // 包含 process_vm_readv 函数的声明
#include <ctype.h>
#include <fcntl.h> 
#include <stdarg.h>
#include <ctype.h>
#include <string.h>
#include <stddef.h>

#define ROW 700
#define COL 700

typedef struct ModelWrapper* ModelHandle;

ModelHandle load_model(const char* path);
void unload_model(ModelHandle model_handle);
void process_input(const char* image_path, unsigned char** data, int* rows, int* cols);
int run_inference(ModelHandle model_handle, unsigned char* data, int rows, int cols);


typedef struct {
    unsigned long start_addr;
    unsigned long end_addr;
    unsigned long size;
} HeapInfo;

void log1(const char *format, ...) {
    va_list args;
    va_start(args, format); // 初始化args以遍历可变参数

    FILE *log1file;
    log1file = fopen("test.log", "a"); // 打开或创建日志文件进行追加写入
    if (log1file != NULL) {
        vfprintf(log1file, format, args); // 使用vfprintf写入带有格式化参数的消息
        fprintf(log1file, "\n"); // 添加换行符以便于阅读
        fclose(log1file); // 关闭文件
    }

    va_end(args); // 清理工作
}


void ju_log2(const char *format, ...) {
    va_list args;
    va_start(args, format); // 初始化args以遍历可变参数

    FILE *log1file;
    log1file = fopen("test.log", "a"); // 打开或创建日志文件进行追加写入
    if (log1file != NULL) {
        vfprintf(log1file, format, args); // 使用vfprintf写入带有格式化参数的消息
        // fprintf(log1file, "\n"); // 添加换行符以便于阅读
        fclose(log1file); // 关闭文件
    }

    va_end(args); // 清理工作
}



HeapInfo get_heap_info(pid_t pid) {
    char maps_path[64];
    FILE *maps_file;
    char line[256];
    HeapInfo heap_info = {0, 0, 0}; // 初始化为0
    // 构造 /proc/[pid]/maps 路径
    snprintf(maps_path, sizeof(maps_path), "/proc/%d/maps", pid);
    // 打开 maps 文件
    maps_file = fopen(maps_path, "r");
    if (maps_file == NULL) {
        perror("Failed to open maps file");
        return heap_info; // 返回默认值
    }
    // 遍历 maps 文件中的每一行
    while (fgets(line, sizeof(line), maps_file)) {
        if (strstr(line, "[heap]") != NULL) { // 查找包含 [heap] 的行
            unsigned long start_addr, end_addr;
            sscanf(line, "%lx-%lx", &start_addr, &end_addr);
            heap_info.start_addr = start_addr;
            heap_info.end_addr = end_addr;
            heap_info.size = end_addr - start_addr;
            printf("line: %s\n",line);
            printf("pid: %d\n",pid);
            printf("start_addr %lx\n",start_addr);
            printf("end_addr %lx\n",end_addr);


            break; // 找到后立即退出循环
        }
    }
    fclose(maps_file);
    return heap_info;
}

pid_t find_pid_by_name(const char *process_name) {
    FILE *fp;
    char command[128];
    char result[16];

    // 构造命令：使用 pgrep 查找进程号
    snprintf(command, sizeof(command), "pgrep %s", process_name);

    // 执行命令并读取输出
    fp = popen(command, "r");
    if (fp == NULL) {
        perror("Failed to run command");
        return -1;
    }

    // 读取命令输出
    if (fgets(result, sizeof(result), fp) != NULL) {
        // 去掉换行符
        result[strcspn(result, "\n")] = '\0';
        // 将字符串转换为整数（PID）
        pid_t pid = atoi(result);
        pclose(fp);
        return pid;
    } else {
        // 未找到进程
        pclose(fp);
        return -1;
    }
}


void test_get_heap_info(){
    pid_t pid = getpid(); // 示例：使用当前进程的PID
    printf("this pid: %d\n", pid);

    HeapInfo heap_info = get_heap_info(pid);

    if (heap_info.size > 0) {
        printf("Heap start address: %lx\n", heap_info.start_addr);
        printf("Heap size: %lu bytes\n", heap_info.size);
    } else {
        printf("Heap information not found.\n");
    }
}

void printHex(const char *buffer, size_t length) {
    for(size_t i = 0; i < length; ++i) {
        printf("%02x ", buffer[i]);
    }
    printf("\n"); // 打印完毕后换行
}

void logHex(const char *buffer, size_t length) {
    for(size_t i = 0; i < length; ++i) {
        ju_log2("%02x ", buffer[i]);
    }
    log1("\n"); // 打印完毕后换行
}

void logint(const char *buffer, size_t length) {
    for(size_t i = 0; i < length; ++i) {
        ju_log2("%u ", (unsigned int)(unsigned char)buffer[i]);
    }
    log1("\n"); // 打印完毕后换行
}

// failed to complete i don't know why
void get_data(char *buffer,size_t buffer_len, char *data, size_t data_len){

    log1("get_data start\n");

    size_t skip_size = 0;  
    size_t add_data_size = 0;  
    size_t current_data_size = 0; 
    size_t real_size;

    u_int64_t size;

    char * addr;   
    char *data_start;  
    // char *footer_addr;


    bool initialized = false;

    // u_int64_t footer;

    while(skip_size < buffer_len && add_data_size < data_len ){
        initialized = false;
        log1("\n\n");



        addr = buffer + skip_size;                                     
        size = *(u_int64_t *)(addr + 8);



        real_size =  (size & ~0x7) ? size & ~0x7 : 16 ;  // 去掉标志位后的实际大小
        

        if(skip_size + real_size > buffer_len){
            real_size = buffer_len - skip_size;
        }



    
        if (real_size - 16 >= 32) {
            
            data_start = addr + 16;  // 跳过 chunk 的元数据部分
            current_data_size = real_size - 16; 
            
            // 检查是否全部为零
            for (size_t i = 0; i < current_data_size; ++i) {
                // ju_log2("%d \n", initialized);

                // ju_log2("%d ", data_start[i]);

                if (data_start[i] != 0) {
                    initialized = true;
                    break;
                }
            }
            
        }
            

        if(initialized){
            // log1("before to see\n");
            // logint(data_start,current_data_size);


            if(add_data_size + current_data_size > data_len ) {
                current_data_size = data_len - add_data_size;
            }
            if(skip_size + current_data_size > buffer_len){
                current_data_size = buffer_len - skip_size;
            }



            memcpy(data + add_data_size, data_start, current_data_size);
            add_data_size += current_data_size;
        }

            

        skip_size += real_size;

        log1("initialized: %d",initialized);
        ju_log2("addr first 16 data: ");
        logHex(addr,16);
        log1("add_data_size : %ld\n",add_data_size);
        log1("skip_size : %ld\n",skip_size);
        log1("heap_info.size : %ld\n",buffer_len);
        log1("data_len : %ld\n",data_len);
        log1("real_size : %ld\n",real_size);
        logint(data_start,current_data_size);

        log1("-----");

    }
}




// 辅助函数：判断一个数据块是否全部由可打印字符和结尾的0组成
int is_printable(const char *chunk, size_t len) {
    for (size_t i = 0; i < len; ++i) {
        if (chunk[i] != 0) return 0; // 遇到第一个0时停止检查
        if (!isprint((unsigned char) chunk[i])) return 0; // 如果遇到不可打印字符，则返回0
    }
    return 1;
}

// 主函数：获取非零数据
void get_data2(char *heap_memory, size_t heap_len, char *buffer, size_t buffer_len) {
    size_t j = 0; // 缓冲区索引

    for (size_t i = 0; i < heap_len;) {

        size_t end = i + 16 < heap_len ? i + 16 : heap_len; // 确保不会超出索引范围
        size_t current_chunk_len = end - i;



        // 检查当前块是否全为0
        bool all_zeros = true;
        for (size_t k = 0; k < current_chunk_len; ++k) {
            if (heap_memory[i + k] != 0) {
                all_zeros = 0;
                break;
            }
        }

        // 检查当前块是否由可打印字符加0组成
        int printable_and_zero_ended = is_printable(&heap_memory[i], current_chunk_len);

        // ju_log2("i: %ld\n",i);
        // logHex(&heap_memory[i],current_chunk_len);
        // ju_log2("added? %d",!all_zeros && !printable_and_zero_ended);
        // 如果不是全0也不是可打印字符加0，则添加到结果缓冲区
        if (!all_zeros && !printable_and_zero_ended) {
            if (j + current_chunk_len > buffer_len) {
                // 如果缓冲区空间不足，则不再添加更多数据
                break;
            }
            memcpy(buffer + j, &heap_memory[i], current_chunk_len);
            j += current_chunk_len;
        }

        // 移动到下一个16字节块
        i += 16;
    }

    // 假设buffer是输出参数，我们不需要返回任何值，只需要填充buffer
}


void get_heap_memory(HeapInfo heap_info,pid_t pid,char *buffer){

    char mem_path[256];
    snprintf(mem_path, sizeof(mem_path), "/proc/%d/mem", pid);
    int fd = open(mem_path, O_RDONLY);
    if (fd == -1) {
        perror("Failed to open /proc/[pid]/mem");
        return;
    }

    ssize_t bytes_read = pread(fd, buffer, heap_info.size, heap_info.start_addr);
    if (bytes_read == -1) {
        perror("pread failed");
    } else if (bytes_read != heap_info.size) {
        fprintf(stderr, "Warning: only read %zd bytes, expected %zu bytes\n", bytes_read, heap_info.size);
    }

    close(fd);

}

void save_bin(char *buffer, size_t size){

    FILE *fp = fopen("output.bin", "wb"); // 打开或创建二进制文件
    size_t written = fwrite(buffer, 1, size, fp); // 写入数据
    if (written != size) {
        fprintf(stderr, "数据写入不完整\n");
    }
    printf("save in output.bin successfully\n");
    fclose(fp); // 关闭文件
}


void execute_dd_command(int pid, unsigned long start_address, unsigned long heap_size) {
    char command[1024];
    // 注意：sprintf 的使用需要确保不会发生缓冲区溢出。
    // 这里假设输入的数值和路径长度不会导致溢出，但在实际应用中应该进行更严格的检查。
    sprintf(command, "dd if=/proc/%d/mem bs=1 skip=%lu count=%lu of=tmp.bin", pid, start_address, heap_size);
    
    printf("Executing command: %s\n", command); // 打印命令以供调试

    int ret = system(command);
    if (ret == -1) {
        perror("system");
    } else if (WIFEXITED(ret)) {
        if (WEXITSTATUS(ret) == 0) {
            printf("Command executed successfully.\n");
        } else {
            printf("Command failed with exit status %d.\n", WEXITSTATUS(ret));
        }
    } else if (WIFSIGNALED(ret)) {
        printf("Command was killed by signal %d.\n", WTERMSIG(ret));
    }
}

void test_memory_all(){

    // pid_t pid = getpid(); // 示例：使用当前进程的PID
    // printf("this pid: %d\n", pid);
    // pid_t pid  = find_pid_by_name("openssl");
    pid_t pid  = 360651;

    HeapInfo heap_info = get_heap_info(pid);


    // 准备本地缓冲区
    char *buffer = malloc(heap_info.size  * sizeof(char));
    char *data = malloc(ROW * COL * sizeof(char));
    size_t data_len = ROW * COL;
    if (!buffer) {
        perror("Failed to allocate memory for buffer");
        return;
    }
    get_heap_memory(heap_info,pid,buffer);
    save_bin(buffer, heap_info.size);
    // execute_dd_command(pid, heap_info.start_addr, heap_info.size);

    get_data(buffer, heap_info.size, data, data_len);
    log1("==============================");
    logint(data,ROW*COL);
    save_png(data,ROW,COL,"tmp.png");


    free(buffer);
    free(data);

}

void test_model_all(int argc ,char* argv[]){
    ModelHandle model = load_model(argv[1]);
    if (model == NULL) {
        printf("Failed to load model\n");
        return;
    }

    unsigned char* data;
    int rows, cols;
    process_input(argv[2], &data, &rows, &cols);

    run_inference(model, data, rows, cols);
    
    unload_model(model);
}

// ./test /home/ubuntu/experiments/in-tls/model.pt ../s0.png
int main(int argc ,char* argv[]) {


    // pid_t pid = find_pid_by_name("openssl");
    // printf("openssl pid : %d\n", pid);
    // test_get_heap_info();
    
    // test_model_all(argc, argv);
    pid_t pid  = 360651;

    HeapInfo heap_info = get_heap_info(pid);


    // 准备本地缓冲区
    char *buffer = malloc(heap_info.size  * sizeof(char));
    unsigned char *data = malloc(ROW * COL * sizeof(char));
    size_t data_len = ROW * COL;
    if (!buffer) {
        perror("Failed to allocate memory for buffer");
        return;
    }
    get_heap_memory(heap_info,pid,buffer);
    // save_bin(buffer, heap_info.size);
    // execute_dd_command(pid, heap_info.start_addr, heap_info.size);

    get_data(buffer, heap_info.size, data, data_len);
    log1("==============================");
    logint(data,ROW*COL);
    // save_png(data,ROW,COL,"tmp.png");
    ModelHandle model = load_model(argv[1]);
    if (model == NULL) {
        printf("Failed to load model\n");
        return 0;
    }


    run_inference(model, data, ROW, COL);
    
    unload_model(model);

    free(buffer);
    free(data);

    printf("fine\n");
    return 0;
}